/*
 * To change this template, choose Tools | Templates
 * and open the template in the editor.
 */
package MAIN;

/**
 *
 * @author anhdung
 */
import DAO.Connect;
import DTO.PacketDTO;
import java.security.KeyPairGenerator;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.KeyFactory;
import java.security.spec.InvalidKeySpecException;
import java.security.NoSuchAlgorithmException;
import java.security.spec.*;
import java.io.*;
import java.math.BigInteger;
import java.security.InvalidKeyException;
import java.security.Key;
import java.security.SecureRandom;
import java.util.logging.Level;
import java.util.logging.Logger;
import javax.crypto.BadPaddingException;
import javax.crypto.Cipher;
import javax.crypto.IllegalBlockSizeException;
import javax.crypto.KeyGenerator;
import javax.crypto.NoSuchPaddingException;
import javax.crypto.SecretKey;
import static DAO.Connect.Connect;
import static DAO.Connect.CloseConnect;
import java.nio.charset.Charset;
import java.security.cert.Certificate;
import java.security.cert.CertificateFactory;
import java.security.cert.X509Certificate;

public class GeneratePublicPrivateKeys {
    
    private Key key;//khong su dung.
    public static String secret_key =  "default";//Session-key (Secret-key)  
    public static boolean use_secret_key = false;
    public void generateKeys() {

        try {

            // Get the public/private key pair
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance("RSA");
            keyGen.initialize(1024);
            KeyPair keyPair = keyGen.genKeyPair();
            PrivateKey privateKey = keyPair.getPrivate();
            PublicKey  publicKey  = keyPair.getPublic();

            // Get the bytes of the public and private keys
            byte[] privateKeyBytes = privateKey.getEncoded();
            byte[] publicKeyBytes  = publicKey.getEncoded();
			
            //converting byte arrays to strings
            String privateString = new String(privateKeyBytes);
            String publicString = new String(publicKeyBytes);
						
            KeyFactory fact = KeyFactory.getInstance("RSA");
            RSAPublicKeySpec pub = fact.getKeySpec(keyPair.getPublic(),RSAPublicKeySpec.class);
            RSAPrivateKeySpec priv = fact.getKeySpec(keyPair.getPrivate(),RSAPrivateKeySpec.class);

            saveToFile("public.key", pub.getModulus(), pub.getPublicExponent());
            saveToFile("private.key", priv.getModulus(), priv.getPrivateExponent());

        } catch (InvalidKeySpecException specException) {

            System.out.println("Exception");
            System.out.println("Invalid Key Spec Exception");
            
        } catch (NoSuchAlgorithmException e) {

            System.out.println("Exception");
            System.out.println("No such algorithm: " + "RSA");

        }
		
        catch(IOException e) {
                System.out.println("IO exception");
        }	

    }
	
    private void generateSessionKey() throws NoSuchAlgorithmException {
        KeyGenerator generator;
	generator = KeyGenerator.getInstance("DES");
	generator.init(new SecureRandom());
        key = generator.generateKey();
    }
    
    private byte[] encrypt(byte[] data) throws IllegalBlockSizeException,
	    BadPaddingException, NoSuchAlgorithmException,
	    NoSuchPaddingException, InvalidKeyException,
	    UnsupportedEncodingException {
	// Get a cipher object.
	Cipher cipher = Cipher.getInstance("DES/ECB/PKCS5Padding");
	cipher.init(Cipher.ENCRYPT_MODE, key);
 
	// encrypt using the cypher
	byte[] raw = cipher.doFinal(data);
 
	return raw;
    }
    
      private byte[] decrypt(byte[] data) throws InvalidKeyException,
	    NoSuchAlgorithmException, NoSuchPaddingException,
	    IllegalBlockSizeException, BadPaddingException, IOException {
 
		// Get a cipher object.
		Cipher cipher = Cipher.getInstance("DES/ECB/PKCS5Padding");
		cipher.init(Cipher.DECRYPT_MODE, key);
 
		//decode the message
		byte[] result = cipher.doFinal(data);
 
		return result;
    }
    
    
    private void saveToFile(String fileName,
      BigInteger mod, BigInteger exp) throws IOException {
      try (ObjectOutputStream oout = new ObjectOutputStream(
                 new BufferedOutputStream(new FileOutputStream(fileName)))) {
            oout.writeObject(mod);
            oout.writeObject(exp);
      } catch (Exception e) {
            throw new IOException("Unexpected error", e);
      }
    }
    

    private PublicKey readPubKeyFromFile(String keyFileName) throws IOException {
      FileInputStream in = new FileInputStream(keyFileName);
      ObjectInputStream oin =
            new ObjectInputStream(new BufferedInputStream(in));
      try {
            BigInteger m = (BigInteger) oin.readObject();
            BigInteger e = (BigInteger) oin.readObject();
            RSAPublicKeySpec keySpec = new RSAPublicKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PublicKey pubKey = fact.generatePublic(keySpec);
            return pubKey;
      } catch (Exception e) {
            throw new RuntimeException("Spurious serialisation error", e);
      } finally {
            oin.close();
      }
    }

    /* THis function implements the reading of the private key from the keyFileName*/
    private PrivateKey readPriKeyFromFile(String keyFileName) throws IOException {

      FileInputStream in = new FileInputStream(keyFileName);
      ObjectInputStream oin =
            new ObjectInputStream(new BufferedInputStream(in));
      try {
            BigInteger m = (BigInteger) oin.readObject();
            BigInteger e = (BigInteger) oin.readObject();
            RSAPrivateKeySpec keySpec = new RSAPrivateKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PrivateKey priKey = fact.generatePrivate(keySpec);
            return priKey;
      } catch (Exception e) {
            throw new RuntimeException("Spurious serialisation error", e);
      } finally {
            oin.close();
      }
    }
//    public static byte[] encryptKey(byte[] data) {
//      try {
//              PublicKey pubKey = readPubKeyFromFile("public.key");
//              Cipher cipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
//              cipher.init(Cipher.ENCRYPT_MODE, pubKey);
//              byte[] cipherData = cipher.doFinal(data);
//              return cipherData;
//       } catch(Exception e) { e.printStackTrace();return null;}
//
//    }
//    /*Decrypting here*/
//    public static byte[] decryptKey(byte[] data) {
//      try {
//              PrivateKey priKey = readPriKeyFromFile("private.key");
//              Cipher dcipher = Cipher.getInstance("RSA/ECB/PKCS1Padding");
//              dcipher.init(Cipher.DECRYPT_MODE, priKey);
//              byte[] cipherData = dcipher.doFinal(data);
//              return cipherData;
//            } catch(Exception e) { e.printStackTrace();return null;}
//
//    }
        
    public boolean sendPubKey() {

        try {
            PublicKey pubKey = readPubKeyFromFile("public.key");
            byte[] pubKeyData = pubKey.getEncoded();
            DTO.PacketDTO pkdto=new DTO.PacketDTO();
            pkdto.setData(pubKeyData);
            Connect.sendData(pkdto);
            
            return true;
        } catch (IOException ex) {
            Logger.getLogger(GeneratePublicPrivateKeys.class.getName()).log(Level.SEVERE, null, ex);
        }

 
        return false;
        
    }
    
    public boolean receivePubKey() {
        Object data = Connect.receiveData();
        byte[] command = new byte[4];
        
        try {
            ObjectInputStream objectInput = new ObjectInputStream((ObjectInputStream) data);
            objectInput.read(command);
                        
            BigInteger m = (BigInteger) objectInput.readObject();
            BigInteger e = (BigInteger) objectInput.readObject();
            RSAPublicKeySpec keySpec = new RSAPublicKeySpec(m, e);
            KeyFactory fact = KeyFactory.getInstance("RSA");
            PublicKey pubKey = fact.generatePublic(keySpec);
            
            saveToFile("public.server.key", m, e);
            return true;
            
        } catch (IOException ex) {
            Logger.getLogger(GeneratePublicPrivateKeys.class.getName()).log(Level.SEVERE, null, ex);
        } catch (ClassNotFoundException ex) {
            Logger.getLogger(GeneratePublicPrivateKeys.class.getName()).log(Level.SEVERE, null, ex);
        } catch (NoSuchAlgorithmException ex) {
            Logger.getLogger(GeneratePublicPrivateKeys.class.getName()).log(Level.SEVERE, null, ex);
        } catch (InvalidKeySpecException ex) {
            Logger.getLogger(GeneratePublicPrivateKeys.class.getName()).log(Level.SEVERE, null, ex);
        }
        
        return false;
    }
    public boolean exchangeKey(PublicKey server_pubkey) {
       try {
           generateSessionKey();

           PrivateKey priKeyClient = readPriKeyFromFile("private.key");
           byte[] priKeyData = priKeyClient.getEncoded();

           //PublicKey pubKeyServer = readPubKeyFromFile("public.server.key");
           PublicKey pubKeyServer = server_pubkey;
           byte[] pubKeyData = pubKeyServer.getEncoded();

           //encypt by private key of client
           Cipher cipherPri = Cipher.getInstance("RSA");//--"RSA/ECB/PKSC1Padding"
           cipherPri.init(Cipher.ENCRYPT_MODE, priKeyClient);
           byte[] cipherDataPri = cipherPri.doFinal(key.getEncoded());
           System.out.println("input length? " + cipherDataPri.length);

           //encrypt by public key of server
           Cipher cipherPub = Cipher.getInstance("RSA");//--"RSA/ECB/PKSC1Padding"
           cipherPub.init(Cipher.ENCRYPT_MODE, pubKeyServer);
           byte[] cipherDataPub = cipherPub.doFinal(cipherDataPri);

            DTO.PacketDTO pkdto=new DTO.PacketDTO();
            pkdto.setData(cipherDataPub);
            pkdto.setCommand(105);//Dung moi them vo
            Connect.sendData(pkdto);

            return true;

       }
       catch (Exception e) {
           e.printStackTrace();
           return false;
       }
    }

    public static void Certificate_And_SecretKey() throws Exception
    {
        //========Phan Chung Thuc=======
        Connect();//gui-nhan lan 1: Certificate va Public Key        
        //**Tao public va private key
        GeneratePublicPrivateKeys gk = new GeneratePublicPrivateKeys();
        gk.generateKeys();
        PrivateKey privk = gk.readPriKeyFromFile("private.key");
        PublicKey pubk = gk.readPubKeyFromFile("public.key");
        //System.out.println(privk);
        //System.out.println(pubk);

        //**Tao Cert tuong ung voi public key
        usc_Certificate.privKey = privk;
        usc_Certificate.pubKey = pubk;
        usc_Certificate.createCertificate("GUEST.cer");

        //**Gui GUEST.cer cho Server
        FileInputStream is = new FileInputStream(new File("GUEST.cer"));
        CertificateFactory cf = CertificateFactory.getInstance("X.509");
        Certificate guest_cer = cf.generateCertificate(is);
        DTO.PacketDTO pdto = new PacketDTO();
        pdto.setCommand(101);//ti nua, ben Server nhan so 101 thi biet la nhan GUEST.cer
        pdto.setData(guest_cer);
        Connect.sendData(pdto);

        //Nhan phan hoi tu Server
        Object obj = Connect.receiveData();
        pdto = new PacketDTO();
        pdto = (PacketDTO)obj;
        if(pdto.getFlag()==false)//Neu Server thay GUEST.cer khong hop le
        {
            System.out.println("Server says: " + pdto.getData().toString());
            CloseConnect();
            return;
            //Close all
        }
        //**Nhan SERVER.cer tu Server            
        System.out.println("== == == == == == == == == ==");
        System.out.println("Server has sent you SERVER.cer ");
        Certificate server_cer = (Certificate)pdto.getData();
        boolean f1 = usc_Certificate.validateCertificate(server_cer);
        if(f1==false)
        {
            System.out.println("You have had SERVER.cer. It is INVALID.");
            //Gui canh bao toi Server va ngat ket noi.
            pdto = new DTO.PacketDTO();
            pdto.setFlag(false);
            pdto.setCommand(102);///ti nua, ben Server nhan so 102 thi biet la SERVER.cer khong hop le
            pdto.setData("Invalid Certificate.");
            Connect.sendData(pdto);
            CloseConnect();
            return;
            //Close all
        }
        System.out.println("You have had SERVER.cer. It is VALID.");
        //Luu SERVER.cer
        FileOutputStream os = new FileOutputStream("SERVER.cer");
        os.write(server_cer.getEncoded());
        os.close();
        
        System.out.println("== == == == == == == == == ==");
        PublicKey server_pubkey = server_cer.getPublicKey();
        
        //
        //Kiem tra Public-Private Key cua Server     
        String plaintext_C =  usc_sessionkey_mac.randomString(50);
        String digest_1 = Hash_SHA256.sha256(plaintext_C);
        byte[] encrypted_digest_1 = usc_RSAKeyPair.PAIR2_Encrypt(
                digest_1.getBytes(Charset.forName("UTF-8")), server_pubkey);
        
        byte[][] data_1_C = new byte[2][];
        data_1_C[0] = plaintext_C.getBytes(Charset.forName("UTF-8"));
        data_1_C[1] = encrypted_digest_1;
        //**Gui chung cho Server
        pdto = new PacketDTO();
        pdto.setCommand(103);//lenh 103: Kiem Tra Public-Private key, yeu cau Secret-key
        pdto.setData((Object)data_1_C);
        System.out.println("You have sent to Server a plaintext: " + plaintext_C);
        System.out.println("And it's SHA-256 digest: " + digest_1);
        System.out.println("-->Use Server PUBLIC KEY to encrypt.");
        Connect();//gui-nhan lan 2: {plaintext + encrypted(hash)} .
        Connect.sendData(pdto);
        System.out.println("-----");
        
        obj = new Object();
        obj = Connect.receiveData();
        pdto = new PacketDTO();
        pdto = (PacketDTO)obj;
        if(pdto.getFlag()==false)
        {
            System.out.println("Server says: " + pdto.getData());
            CloseConnect();
            return;
        }
        //Kiem tra Public-Private Key cua Client, kem nhan luon secret-key (session-key)
        byte[][] data_2_S = (byte[][])pdto.getData();
        byte[] encrypted_plaintext_S = data_2_S[0];
        byte[] decrypted_plaintext_S = usc_RSAKeyPair.PAIR2_Decrypt(encrypted_plaintext_S, usc_Certificate.privKey);
        //String plaintext_S = new String(data_2_S[0],"UTF-8");
        String plaintext_S = new String(decrypted_plaintext_S,"UTF-8");
        System.out.println("Server has sent a plaintext: " + plaintext_S);
        
        byte[] encrypted_digest_2 = data_2_S[1];        
        byte[] decrypted_digest_2 = usc_RSAKeyPair.PAIR2_Decrypt(encrypted_digest_2, usc_Certificate.privKey);
        System.out.println("-->Use Client PRIVATE KEY to decrypt.");
        String digest_2 = new String(decrypted_digest_2,"UTF-8");
        System.out.println("With SHA-256 Digest: " + digest_2);
        System.out.print("Comparing SHA-256 Digest: ");
        String hash_result = Hash_SHA256.sha256(plaintext_S);
        if(hash_result.equals(digest_2)==false)
        {
            Connect();
            System.out.println("NOT CORRECT! That is not my server. Disconect...");
            pdto = new PacketDTO();
            pdto.setCommand(104);
            pdto.setData("Server Key Pair is not valid. Disconect!");
            Connect.sendData(pdto);
            CloseConnect();
            return;
        }
        System.out.println("CORRECT! That is my SERVER.");
        //~~~~~SECRET-KEY
        secret_key = usc_sessionkey_mac.xorMessage(plaintext_S, plaintext_C) ;
        System.out.println("SECRET KEY: " + secret_key);
        use_secret_key=true;
        //Bao lai cho Server da nhan duoc Secret key
        Connect();
        pdto = new PacketDTO();
        pdto.setCommand(105);
        Connect.sendData(pdto);
        System.out.println();
        System.out.println("-> Send reponse to Server about receiving Secret Key.");
        System.out.println("Now. use_secret_key is ON.");
        System.out.println("== == == == == == == == == ==");  
        
        CloseConnect();
        /*
        //Su dung exchangeKey that bai, vi input.length=128 trong khi gioi han input RSA la 117
        GeneratePublicPrivateKeys g = new GeneratePublicPrivateKeys();
        g.exchangeKey(server_pubkey);
        */          
    }
    
    public static void main(String[] args) throws Exception
    {
        //**Trao doi Certificate, tao Secret Key (Session Key)
        //Certificate_And_SecretKey();        
        
        /*
        //**Su dung Secret Key de trao doi, co kem MAC = HASH{plaintext, secret key}
        //System.out.println("secret_key: " + secret_key);
        String data = "Lenh dat ve.";
        byte[] macBytes = null;
        macBytes = usc_sessionkey_mac.generateMAC(data, secret_key);        
        String MAC = new String(macBytes,"UTF-8");
        byte[] macBytesAfter = MAC.getBytes(Charset.forName("UTF-8"));
        if(macBytesAfter.equals(macBytes)==true)
        {
            System.out.println("2 MAC is equal.");
        }
        else
        {
            System.out.println("2 MAC is not equal. You shoudn't use String for MAC");
        }
        
        DTO.PacketDTO pdto = new PacketDTO();
        pdto.setData(null);        
        pdto.setCommand(106);//ti nua, ben Server nhan so 106 thi biet la test Secret-key va MAC.
        Connect();
        Connect.sendData(pdto);
        */
        
        /*
        //Test Secret-key: Command 106
        Connect();
        String s = "I am talent.";
        PacketDTO pdto = new PacketDTO();
        pdto.setCommand(106);
        pdto.setData(s);
        Connect.sendData(pdto);
        System.out.println("Send: " + s);
        
        Object reobj = Connect.receiveData();
        pdto = (PacketDTO)reobj;
        String s2 = pdto.getData().toString();
        System.out.println("Receive: " + s2);        
        CloseConnect();
        */
        //System.out.println("END.");
        
    }
    
}
